package com.fuyun.rpc.client.loadbalancing;

import com.fuyun.rpc.common.domain.ServiceMetaInfo;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 加权轮询负载均衡器
 */
public class WeightedRoundRobinLoadBalancer implements LoadBalancer {

	/**
	 * 服务权重映射
	 * 在实际应用中，这个权重可以从配置文件或注册中心获取
	 */
	private final Map<String, Integer> serviceWeights = new ConcurrentHashMap<>();

	/**
	 * 当前权重计数器
	 */
	private final Map<String, AtomicInteger> currentWeights = new ConcurrentHashMap<>();

	public WeightedRoundRobinLoadBalancer() {
		// 示例权重配置，实际应用中应该从配置文件读取
		serviceWeights.put("192.168.1.1", 5);  // 权重5
		serviceWeights.put("192.168.1.2", 3);  // 权重3
		serviceWeights.put("192.168.1.3", 2);  // 权重2
	}

	@Override
	public ServiceMetaInfo select(Map<String, Object> requestParams, List<ServiceMetaInfo> serviceMetaInfoList) {
		if (serviceMetaInfoList == null || serviceMetaInfoList.isEmpty()) {
			return null;
		}

		if (serviceMetaInfoList.size() == 1) {
			return serviceMetaInfoList.get(0);
		}

		// 查找最高权重
		int maxWeight = 0;
		for (ServiceMetaInfo serviceMetaInfo : serviceMetaInfoList) {
			String host = serviceMetaInfo.getServiceHost();
			int weight = serviceWeights.getOrDefault(host, 1);
			maxWeight = Math.max(maxWeight, weight);
		}

		// 执行加权轮询算法
		while (true) {
			ServiceMetaInfo selectedService = null;
			int maxCurrentWeight = -1;

			for (ServiceMetaInfo serviceMetaInfo : serviceMetaInfoList) {
				String host = serviceMetaInfo.getServiceHost();
				int weight = serviceWeights.getOrDefault(host, 1);

				// 更新当前权重
				AtomicInteger currentWeight = currentWeights.computeIfAbsent(host, k -> new AtomicInteger(0));
				currentWeight.addAndGet(weight);

				// 查找当前权重最大的服务
				if (currentWeight.get() > maxCurrentWeight) {
					maxCurrentWeight = currentWeight.get();
					selectedService = serviceMetaInfo;
				}
			}

			if (selectedService != null) {
				String selectedHost = selectedService.getServiceHost();
				// 将选中服务的当前权重减去最大权重
				currentWeights.get(selectedHost)
					.addAndGet(-maxWeight);
				return selectedService;
			}

			// 防止无限循环
			break;
		}

		// 如果算法未能选择服务，则回退到普通轮询
		return new RoundRobinLoadBalancer().select(requestParams, serviceMetaInfoList);
	}

	/**
	 * 设置服务权重
	 * @param host 服务主机地址
	 * @param weight 权重值
	 */
	public void setServiceWeight(String host, int weight) {
		if (weight > 0) {
			serviceWeights.put(host, weight);
		}
	}
}